Parallelize and right-size the sequence-properties compute step#22
Parallelize and right-size the sequence-properties compute step#22PaulNewling wants to merge 6 commits into
Conversation
Scale the compute step with the dataset while keeping output byte-identical (the quantized-equal / CID-stability contract). Deterministic threading: the workflow requests cores and passes them to polars via POLARS_MAX_THREADS; main.py pins BLAS/OpenMP intra-op threads to 1 so the order-sensitive 'counts @ weights' reduction never splits across threads. Emitted bytes are independent of core count. Input-scaled resources: replace the flat mem/cpu with exec resource formulas — a lineCount-based memFormula for both modes (fallback 12GiB), and a cpuFormula up to 4 cores for peptide/amplicon (antibody stays single-core; it is ~1.1x on threads). Constants calibrated from measured peak RSS (~3 KiB/row peptide, ~4.3 KiB/row antibody) and thread scaling. Compute engine: _charge_raw computes 10**ph once and reuses it via scalar factors (~2.4x on the charge/pI path); peptide cleans its column once and shares the intermediate between counts and instability; count/instability scatters use np.bincount instead of np.add.at (~3x). aa_fraction long-format uses a polars unpivot instead of a 20xN Python-list explosion. Peptide compute ~2.3x, antibody wall ~25% at 1M rows. Verified: 261 software tests pass; output byte-identical vs the pre-change implementation for peptide and antibody (incl. empty/stop-codon/lowercase/non-standard/missing-region edge cases) and across 1 vs 4 polars threads. Bumps @platforma-sdk/workflow-tengo to 6.6.5 (catalog) for the resource-formula API.
…stat Second round of compute-engine speedups on top of the previous commit; output stays byte-identical (quantized-equal / CID-stability contract). Vectorize the per-chain median-CDR3-length stat (_median_cdr3_length_by_chain) with a polars count_matches over standard-AA chars instead of one effective_length Python call per row — the largest antibody-mode stats cost. Median formula unchanged, so stats.json is byte-identical. Row-parallelize the pI bisection (isoelectric_point / fv_isoelectric_point): contiguous, data-independent row blocks bisected on separate threads (numpy elementwise releases the GIL), concatenated in row order. Each row's bisection depends only on its own counts, so the result is bit-identical regardless of worker count. With both, antibody now scales with cores (~1.13x -> ~1.53x on 4 cores; 10.0s -> 4.6s at 1M rows), so the workflow scales antibody CPU up to 4 as well (it was pinned to cpu(1) when the mode did not parallelize). Both modes now share the same cpuFormula. Add a determinism test above vectorized._PARALLEL_MIN_ROWS (antibody, exercising both isoelectric_point and fv_isoelectric_point) asserting byte-identical output at 1 vs 4 threads; the smaller cases stay under the parallel threshold and cannot reach it. Verified: 262 software tests pass; output byte-identical vs the pre-change implementation (peptide + antibody at 1M) and across 1 vs 4 threads at scale.
…equence OOM
The lineCount-only memFormula under-provisioned low-row/long-sequence inputs: peak RSS is dominated by O(total-residues) cleaning transients (the DIWV weight gather, residue-row ids, byte buffers), which lineCount does not see. An amplicon-like 200k-row x ~600aa input peaks at 4.51 GiB but the old formula sized it to 200k*5KiB -> the 2 GiB floor, an OOM the pre-change flat 12 GiB avoided.
Add a residue-scaled term size("input.tsv") * 64 (measured ~40 B peak RSS per input byte; 64x for headroom) and take max() with the row-scaled lineCount term, so each dominates in its regime: short sequences use the row term (no change), long sequences use the size term. Verified: amplicon case now sizes to ~7.25 GiB >= 4.51 GiB peak; normal peptide/antibody sizing is unchanged (row term still dominates).
Review follow-ups #3 and #4, both byte-identical (no numeric or output change). #3: _charge_raw takes an optional precomputed ten_ph; fv_isoelectric_point computes 10**ph once and shares it across both chains instead of recomputing the array-wide transcendental twice per Fv bisection step. Same value -> bit-identical. #4: build the median-CDR3 count_matches character class from STANDARD_AAS (_STANDARD_AA_CLASS) instead of a hardcoded 40-char literal, so it cannot drift from properties.effective_length / the count substrate. Same class -> bit-identical. Also softened the _charge_raw docstring's absolute 'bytes/CID unchanged' to 'within the quantized-equal contract' (review #5): the FP drift is byte-identical on the corpus but a knife-edge bisection midpoint within ~1e-15 of zero is a rare theoretical exception. Verified: 262 tests pass, ruff clean, output byte-identical vs prior HEAD on antibody 1M (exercises both the Fv pI path and the median-CDR3 stat).
…t the static cap Resolves review #2 (thread oversubscription). The workflow pinned POLARS_MAX_THREADS to the static cap (computeCpu=4) while cpuFormula grants 1..4 cores dynamically, and vectorized._n_workers() reads POLARS_MAX_THREADS to size the pI-bisection pool. So a sub-cap input (>=50k but <250k rows -> 1 core granted) spawned 4 polars + 4 bisection threads on a 1-core quota. Set it instead with envExpr("POLARS_MAX_THREADS", "{system.cpu}"): the backend's command-expression engine substitutes the resolved quota grant (ComputeLimitsInfo.CPU = the cpuFormula result) before the process starts, so main.py reads the true granted core count at import. Gated on feats.commandExpressions with a static env(string(computeCpu)) fallback for older backends. This also removes the leaky reuse of the static cap and keeps the thread count in lockstep with the CPU reservation. Verified: workflow tengo compiles (feats/envExpr), 262 tests pass. The runtime {system.cpu} substitution resolves at execution time, so it needs the live-backend integration test to exercise end-to-end.
Changeset: split the memory/CPU/threads wall into one-idea paragraphs (memory, CPU, determinism), convert the compute-optimization run-in list to scannable bullets, cut wordiness, and drop repeated byte-identity claims.
Comments: make the aa_fraction build comment active voice ('is unpivoted'/'are made contiguous' -> active), and fix a stale _n_workers docstring (POLARS_MAX_THREADS is the granted cores now, not the static cpu() request). No code changes; ruff + tests green.
There was a problem hiding this comment.
Code Review
This pull request introduces parallelization and resource scaling to the compute step, optimizing performance while maintaining strict output determinism. Key optimizations include cleaning columns once, replacing slow np.add.at scatters with np.bincount, vectorizing CDR3 length calculations, and restructuring the peptide amino acid fraction calculation using a Polars unpivot. To guarantee byte-identical outputs, BLAS/OpenMP intra-op threads are pinned to 1. Feedback on these changes highlights opportunities to improve performance and type safety: removing a redundant np.ascontiguousarray call in the peptide path, correcting incomplete type annotations for ph and ten_ph in _charge_raw, and adding missing type annotations to the new _bisect_rows_parallel helper.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| wide = pl.DataFrame( | ||
| {"entity_key": key_series} | ||
| | {aa: _na_to_null(np.ascontiguousarray(fractions[:, i])) for i, aa in enumerate(STANDARD_AAS)} | ||
| ) |
There was a problem hiding this comment.
The use of np.ascontiguousarray here is redundant. Since _na_to_null iterates over the elements of the column slice to convert NaN values to None (resulting in a Python list or object array), the memory contiguity of the input numpy array does not provide any performance benefit. Calling np.ascontiguousarray creates an unnecessary intermediate copy of the float array slice for each of the 20 columns, which wastes CPU cycles and memory allocations.
| wide = pl.DataFrame( | |
| {"entity_key": key_series} | |
| | {aa: _na_to_null(np.ascontiguousarray(fractions[:, i])) for i, aa in enumerate(STANDARD_AAS)} | |
| ) | |
| wide = pl.DataFrame( | |
| {"entity_key": key_series} | |
| | {aa: _na_to_null(fractions[:, i]) for i, aa in enumerate(STANDARD_AAS)} | |
| ) |
| def _charge_raw( | ||
| sub: Substrate, ph: float, pka_set: PKaSet, include_cys: bool, ten_ph: np.ndarray | None = None | ||
| ) -> np.ndarray: |
There was a problem hiding this comment.
The type annotations for ph and ten_ph are incomplete and can cause type-checker (e.g., MyPy) errors:
phis annotated asfloat, but during the vectorized bisection, it is passed as a numpy array (np.ndarray).ten_phis annotated asnp.ndarray | None, but ifphis a scalarfloatandten_phis not provided,ten_phis computed as10.0**ph(which is afloat). Assigning afloatto a variable annotated asnp.ndarray | Nonewill trigger a type-checker warning.
Updating the annotations to support both scalar and array types resolves these issues.
def _charge_raw(
sub: Substrate,
ph: float | np.ndarray,
pka_set: PKaSet,
include_cys: bool,
ten_ph: float | np.ndarray | None = None,
) -> np.ndarray:| return 1 | ||
|
|
||
|
|
||
| def _bisect_rows_parallel(make_charge_fn, valid, lo=0.0, hi=14.0, tol=1e-3): |
There was a problem hiding this comment.
The helper function _bisect_rows_parallel is missing type annotations, whereas the rest of the file consistently uses type annotations. Adding type annotations improves code readability, maintainability, and type safety.
Note: This requires importing Callable from typing (i.e., from typing import Callable).
def _bisect_rows_parallel(
make_charge_fn: Callable[[int, int], Callable[[float | np.ndarray], np.ndarray]],
valid: np.ndarray,
lo: float = 0.0,
hi: float = 14.0,
tol: float = 1e-3,
) -> np.ndarray:
Summary
The compute step ran single-threaded under a flat
mem("12GiB")/cpu(1), leaving cores idle and risking OOM on large datasets. This PR parallelizes it, sizes its memory and CPU to the input, and speeds up the compute engine — all while keeping output byte-identical (the quantized-equal / CID-stability contract).At 1M rows: peptide compute ~2.3× faster; antibody wall time ~10.0s → ~4.6s (it now scales with cores, 1.1× → 1.5× on 4).
What changed
Deterministic multithreading.
main.pypins the BLAS/OpenMP intra-op threads to 1, so the one order-sensitive step (thecounts @ weightsreduction) never splits across threads. Polars then parallelizes only order-stable work (CSV read, chain reconstruction, the aa_fraction reshape, the unique-key sort, CSV write). The emitted bytes are independent of core count.Input-scaled resources (replacing the flat mem/cpu):
memFormula(clamp(max(lineCount·perRow, size·64), 2GiB, 64GiB))— a row-scaled term for per-clone structures (count matrix, output columns, aa_fraction) plus a residue-scaledsizeterm for the O(residues) cleaning transients. The size term keeps long-sequence/low-row inputs (amplicon variants) off the 2 GiB floor. Constants are calibrated from measured peak RSS.cpuFormulascales both modes up to 4 cores.POLARS_MAX_THREADSis set to the cores the backend actually grants, via the{system.cpu}command expression, so a sub-cap allocation never oversubscribes its quota.fallback: "12GiB", staticcpu/env).Compute-engine optimizations (all output-preserving):
_charge_rawcomputes10**phonce and reuses it via scalar factors (~2.4× on the charge/pI path).np.bincountinstead of the slowernp.add.at(~3×, bit-identical).count_matchesinstead of one Python call per row.aa_fractionuses a polars unpivot instead of a 20×N Python-list explosion, cutting peak memory.Determinism
Output stays within the quantized-equal contract, verified byte-for-byte:
The charge refactor drifts ~1e-15, below the 3-dp quantization. The one theoretical exception is a bisection midpoint whose net charge lands within ~1e-15 of zero (astronomically rare, never observed).
Testing
ruff check+ruff formatclean.pnpm build:devgreen (7/7 tasks; the workflow Tengo compiles).Not yet verified
The resource formulas and
{system.cpu}resolve at execution time, so the runtime path (thelineCountpre-exec → computed quota, and the command-expression substitution) needs a live-backend run to exercise end-to-end. The Tengo compiles and the mechanism is confirmed against the backend source, but it has not yet run against a live server.